
#include "heaviside_tiling.h"
#include "register/op_def_registry.h"

constexpr uint32_t BLOCK_SIZE = 512;
namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  HeavisideTilingData tiling;
//   const gert::StorageShape* x1_shape = context->GetInputShape(0);
//   int32_t data_sz = 1;
//   for (int i = 0; i < x1_shape->GetStorageShape().GetDimNum(); i++)
//     data_sz *= x1_shape->GetStorageShape().GetDim(i);
  
  
  //获取输入shape信息
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize(); //输入数量
    uint32_t inputBytes = GetSizeByDataType(context->GetInputDesc(0)->GetDataType()); //输入类型
    uint32_t inputLength = inputBytes * inputNum; //输入长度

    uint32_t valuesNum = context->GetInputShape(1)->GetStorageShape().GetShapeSize(); //输入数量

    uint32_t BLOCK_DIM = 40;
    if(valuesNum == 1)
    {
        BLOCK_DIM = 12;
        if(inputLength < 12*256)
            BLOCK_DIM = 1;
        
        uint32_t ALIGN_NUM = BLOCK_SIZE / inputBytes;
        uint32_t totalLengthAligned = ((inputNum + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % BLOCK_DIM;
        uint32_t tailNum = BLOCK_DIM - formerNum;
        uint32_t formerLength = ((totalLengthAligned / BLOCK_DIM + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / ALIGN_NUM) * ALIGN_NUM;

        tiling.set_formerNum(formerNum);
        tiling.set_tailNum(tailNum);
        tiling.set_formerLength(formerLength);
        tiling.set_tailLength(tailLength);
        context->SetTilingKey(1);
    }
    else if(valuesNum == inputNum)
    {
        if(inputLength < 40*256)
            BLOCK_DIM = 1;
        
        uint32_t ALIGN_NUM = BLOCK_SIZE / inputBytes;
        uint32_t totalLengthAligned = ((inputNum + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        uint32_t formerNum = (totalLengthAligned / ALIGN_NUM) % BLOCK_DIM;
        uint32_t tailNum = BLOCK_DIM - formerNum;
        uint32_t formerLength = ((totalLengthAligned / BLOCK_DIM + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
        uint32_t tailLength = (totalLengthAligned / BLOCK_DIM / ALIGN_NUM) * ALIGN_NUM;

        tiling.set_formerNum(formerNum);
        tiling.set_tailNum(tailNum);
        tiling.set_formerLength(formerLength);
        tiling.set_tailLength(tailLength);
        context->SetTilingKey(2);
    }
    else
    {
        BLOCK_DIM = 1;
        const gert::StorageShape* input_StorageShape = context->GetInputShape(0);
        const gert::StorageShape* values_StorageShape = context->GetInputShape(1);
    
        uint32_t input_size = input_StorageShape->GetStorageShape().GetShapeSize();
        uint32_t values_size = values_StorageShape->GetStorageShape().GetShapeSize();

        uint32_t input_dim = input_StorageShape->GetStorageShape().GetDimNum();
        uint32_t values_dim = values_StorageShape->GetStorageShape().GetDimNum();
        uint32_t out_size;
        uint32_t out_dim;
        uint32_t input_shape[8];
        uint32_t values_shape[8];
        uint32_t out_shape[8];
        out_dim = input_dim;
        if(out_dim < values_dim) out_dim = values_dim;

        // out_size = input_size;
        // if(out_size < values_size) out_size = values_size;

        uint32_t input_start = out_dim - input_dim;
        uint32_t values_start = out_dim - values_dim;

        out_size = 1;

        for(int i=0; i< out_dim; i++)
        {
            if(i < input_start) input_shape[i] = 1;
            else input_shape[i] = input_StorageShape->GetStorageShape().GetDim(i - input_start);

            if(i < values_start) values_shape[i] = 1;
            else values_shape[i] = values_StorageShape->GetStorageShape().GetDim(i - values_start);


            out_shape[i] = input_shape[i];
            if(out_shape[i] < values_shape[i]) out_shape[i] = values_shape[i];

            out_size *= out_shape[i];
        }

        tiling.set_out_size(out_size);
        tiling.set_out_dim(out_dim);
        tiling.set_input_shape(input_shape);
        tiling.set_values_shape(values_shape);
        tiling.set_out_shape(out_shape);

        context->SetTilingKey(3);
    }

    context->SetBlockDim(BLOCK_DIM);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Heaviside : public OpDef {
public:
    explicit Heaviside(const char* name) : OpDef(name)
    {
        this->Input("input")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("values")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("out")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(Heaviside);
}
